"""
这个目的是为了让大模型给一段数据查询需求做标记和打分；
"""
from nl2sql.prompt.parse_query_prompt import parser_query_prompt
import ast
import json
from openai import OpenAI
from nl2sql.model.llm import LLM


class ParseQuery:
    def __init__(self,
                 model: str = None,
                 api_key: str = None,
                 base_url: str = None,
                 temperature: float = 0) -> None:
        """
        ParserQuery的初始化函数
        :param model: 模型的名称
        :param api_key:
        :param base_url:
        """

        self.model = model or "qwen-turbo"
        self.api_key = api_key or "sk-34f5c792513c423f90b404d28b070f1f"
        self.base_url = base_url or "https://dashscope.aliyuncs.com/compatible-mode/v1"

        self.llm = LLM(api_key=self.api_key,
                       base_url=self.base_url,
                       model_name=self.model,
                       temperature=temperature)

    def parse_metadata(self, query: str):

        response = self.llm.chat(parser_query_prompt.format(query=query))
        print("Raw model response:\n", response)

        try:
            # 去除 markdown 包裹
            response = response.strip()
            if response.startswith("```"):
                response = response.strip("`").strip("python").strip()

            parsed_response = json.loads(response.replace("'", '"'))  # 单引号转双引号
            if (
                    isinstance(parsed_response, dict) and
                    'tags' in parsed_response and
                    'rating' in parsed_response
            ):
                return parsed_response
            else:
                return None
        except Exception as e:
            print("解析失败:", e)
            return None


if __name__ == '__main__':
    p = ParseQuery()
    res = p.parse_metadata(query='查询近一年以来每天的销量平均值是多少')
    print(res)
